#!/bin/bash

#修改vllm_config.yaml中的ssh配置
CONFIG_FILE="./vllm_config.yaml"
source ../src/read_config.sh

log_dir=${LOGSDIR:-"./qwen3_logs"}
log_file="$log_dir/ssh_vllm.log"
user=$USER
image=$IMAGE
mount_dir=$MOUNT_DIR
container_dir=$CONTAINER_DIR
container_name=$CONTAINER_NAME
mindie_ascend_deploy_dir=$MINDIE_ASCEND_DEPLOY_DIR
delete_container=$DELETE_CONTAINER
other_commands=$OTHER_COMMANDS

check_ssh_key() {
    local ip=$1
    ssh -o BatchMode=yes -o ConnectTimeout=5 $ip echo "SSH_OK" &> /dev/null
    if [ $? -ne 0 ]; then
        echo "Error: SSH key not set up for $ip. please use ssh_rank_table.py or manually configure the ip."
        exit 1
    fi
}

execute_container_commands() {
    local IP=$1
    local MIES_CONTAINER_IP=$2
    local USE_SSH=false

    if [ "$IP" == "$(echo "$IPS" | head -n 1 | awk '{print $1}')" ]; then
        echo "Processing master node at IP: $IP with container IP: $MIES_CONTAINER_IP"
        CONTAINER_EXISTS=$(docker ps -aqf name=$container_name)
    else
        echo "Processing node at IP: $IP with container IP: $MIES_CONTAINER_IP"
        check_ssh_key $IP
        USE_SSH=true
        CONTAINER_EXISTS=$(ssh $user@$IP "docker ps -aqf name=$container_name")
    fi

    if [ ! -z "$CONTAINER_EXISTS" ] && [ "$delete_container" = true ]; then
        echo "Container $CONTAINER_EXISTS already exists. Deleting it..."
        if $USE_SSH; then
            ssh $user@$IP "docker rm -f $container_name"
        else
            docker rm -f $container_name
        fi
        CONTAINER_EXISTS=""
    fi

    if [ -z "$CONTAINER_EXISTS" ]; then
        echo "Container does not exist. Creating a new one..."
        docker_command="docker run -itd -u 0 -e HCCL_IF_IP=$IP --ipc=host --network host \
            --name $container_name \
            --privileged \
            --device=/dev/davinci_manager \
            --device=/dev/devmm_svm \
            --device=/dev/hisi_hdc \
            -v /usr/local/Ascend/driver:/usr/local/Ascend/driver:ro \
            -v /usr/local/sbin:/usr/local/sbin:ro \
            -v $mount_dir:$container_dir \
            $image \
            /bin/bash"

        if $USE_SSH; then
            CONTAINER_ID=$(ssh $user@$IP "$docker_command")
        else
            CONTAINER_ID=$(eval "$docker_command")
        fi
        if [ ! -z "$other_commands" ]; then exec_in_container "$other_commands"; fi
    else
        CONTAINER_ID=$CONTAINER_EXISTS
    fi
    exec_in_container "cd $container_dir/$mindie_ascend_deploy_dir/vllm && mkdir -p $log_dir && bash ./start_vllm.sh >> $log_file 2>&1 &"
}

exec_in_container() {
    local command=$1
    if $USE_SSH; then
        ssh $user@$IP "docker exec $CONTAINER_ID bash -c $command "
    else
        docker exec $CONTAINER_ID bash -c "${command//\'/}"
    fi
}

while read -r IP MIES_CONTAINER_IP; do
    execute_container_commands "$IP" "$MIES_CONTAINER_IP"
done <<< "$IPS"
echo "查看日志: $log_file"
tail -f $log_file